import torch
import math
from app.schemas.action import AttachFileDTO
from app.services.detect.rec_plate import build_detect_model, build_rec_model, run_inference as detect_car_plates


device = None
model_car_plate = None
model_rec_color = None


def load_yolo_model():
    global device
    global model_car_plate
    global model_rec_color
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model_car_plate = build_detect_model('assets/car_plate.pt', device)
    model_rec_color = build_rec_model('assets/plate_rec_color.pth', device)

def predict_car_plate(request: AttachFileDTO):
    return detect_car_plates(model_car_plate, model_rec_color, device, request)


def calc_imgsz(size, stride):
    return math.ceil(size / stride) * stride